using System;
using System.Collections;
using System.Collections.Generic;
using System.Diagnostics;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;

namespace HTCollections.Unsafe
{
    [StructLayout(LayoutKind.Sequential)]
    [DebuggerDisplay("Length = {Length}, Capacity = {Capacity}, IsCreated = {IsCreated}")]
    [BurstCompatible(GenericTypeArguments = new[] { typeof(int), typeof(int) })]
    public unsafe struct UnsafeHashMap<TKey, TValue> : IEnumerable<KeyValuePair<TKey, TValue>>, IDisposable
        where TKey : unmanaged, IEquatable<TKey>
        where TValue : unmanaged

    {
        [NativeDisableUnsafePtrRestriction] public TValue* DataPtr;
        [NativeDisableUnsafePtrRestriction] public int* BucketsPtr;
        [NativeDisableUnsafePtrRestriction] public UnsafeSlot* SlotsPtr;

        public int Capacity;
        public int BucketSize;
        public int Length;
        public int FreeIndex; //空位索引
        public int LastIndex; //顺序往后的索引
        public bool IsCreated;
        public Allocator Allocator;

        int mValueSize;
        int mIntSize;
        int mSlotSize;

        public static UnsafeHashMap<TKey, TValue> Create(int capacity = 4, Allocator allocator = Allocator.Persistent)
        {
            UnsafeHashMap<TKey, TValue> hashMap = new UnsafeHashMap<TKey, TValue>(capacity, allocator);
            return hashMap;
        }

        public UnsafeHashMap(Allocator allocator = Allocator.Persistent)
        {
            Allocator = allocator;
            Capacity = 4;
            BucketSize = 4;
            mIntSize = sizeof(int);
            mValueSize = sizeof(TValue);
            mSlotSize = sizeof(UnsafeSlot);
            DataPtr = UnsafeUtil.Malloc<TValue>(Capacity, mValueSize, Allocator);
            UnsafeUtility.MemClear(DataPtr, Capacity * mValueSize);
            BucketsPtr = UnsafeUtil.Malloc<int>(BucketSize, mIntSize, Allocator);
            UnsafeUtility.MemClear(BucketsPtr, BucketSize * mIntSize);
            SlotsPtr = UnsafeUtil.Malloc<UnsafeSlot>(Capacity, mSlotSize, Allocator);
            UnsafeUtility.MemClear(SlotsPtr, Capacity * mSlotSize);
            Length = 0;
            FreeIndex = -1;
            LastIndex = 0;
            IsCreated = true;
        }

        public UnsafeHashMap(int capacity, Allocator allocator = Allocator.Persistent)
        {
            Allocator = allocator;
            Capacity = capacity;
            BucketSize = capacity;
            mIntSize = sizeof(int);
            mValueSize = sizeof(TValue);
            mSlotSize = sizeof(UnsafeSlot);
            DataPtr = UnsafeUtil.Malloc<TValue>(Capacity, mValueSize, Allocator);
            UnsafeUtility.MemClear(DataPtr, Capacity * mValueSize);
            BucketsPtr = UnsafeUtil.Malloc<int>(BucketSize, mIntSize, Allocator);
            UnsafeUtility.MemClear(BucketsPtr, BucketSize * mIntSize);
            SlotsPtr = UnsafeUtil.Malloc<UnsafeSlot>(Capacity, mSlotSize, Allocator);
            UnsafeUtility.MemClear(SlotsPtr, Capacity * mSlotSize);
            Length = 0;
            FreeIndex = -1;
            LastIndex = 0;
            IsCreated = true;
        }

        public UnsafeHashMap(int capacity, int bucketSize, Allocator allocator = Allocator.Persistent)
        {
            Allocator = allocator;
            Capacity = capacity;
            BucketSize = bucketSize;
            mIntSize = sizeof(int);
            mValueSize = sizeof(TValue);
            mSlotSize = sizeof(UnsafeSlot);
            DataPtr = UnsafeUtil.Malloc<TValue>(Capacity, mValueSize, Allocator);
            UnsafeUtility.MemClear(DataPtr, Capacity * mValueSize);
            BucketsPtr = UnsafeUtil.Malloc<int>(BucketSize, mIntSize, Allocator);
            UnsafeUtility.MemClear(BucketsPtr, BucketSize * mIntSize);
            SlotsPtr = UnsafeUtil.Malloc<UnsafeSlot>(Capacity, mSlotSize, Allocator);
            UnsafeUtility.MemClear(SlotsPtr, Capacity * mSlotSize);
            Length = 0;
            FreeIndex = -1;
            LastIndex = 0;
            IsCreated = true;
        }

        public UnsafeHashMap(UnsafeHashMap<TKey, TValue> hashMap, Allocator allocator = Allocator.Persistent)
        {
            Allocator = allocator;
            Capacity = hashMap.Capacity;
            BucketSize = hashMap.BucketSize;
            mIntSize = sizeof(int);
            mValueSize = sizeof(TValue);
            mSlotSize = sizeof(UnsafeSlot);
            DataPtr = UnsafeUtil.Malloc<TValue>(Capacity, mValueSize, Allocator);
            UnsafeUtility.MemClear(DataPtr, Capacity * mValueSize);
            BucketsPtr = UnsafeUtil.Malloc<int>(BucketSize, mIntSize, Allocator);
            UnsafeUtility.MemClear(BucketsPtr, BucketSize * mIntSize);
            SlotsPtr = UnsafeUtil.Malloc<UnsafeSlot>(Capacity, mSlotSize, Allocator);
            UnsafeUtility.MemClear(SlotsPtr, Capacity * mSlotSize);
            Length = hashMap.Length;
            FreeIndex = hashMap.FreeIndex;
            LastIndex = hashMap.LastIndex;
            UnsafeUtility.MemCpy(hashMap.DataPtr, DataPtr, Length * mValueSize);
            UnsafeUtility.MemCpy(hashMap.BucketsPtr, BucketsPtr, BucketSize * mIntSize);
            UnsafeUtility.MemCpy(hashMap.SlotsPtr, SlotsPtr, Length * mSlotSize);
            IsCreated = true;
        }

        public void Dispose()
        {
            if (!IsCreated) return;
            UnsafeUtil.Free(DataPtr, Allocator);
            DataPtr = null;
            UnsafeUtil.Free(BucketsPtr, Allocator);
            BucketsPtr = null;
            UnsafeUtil.Free(SlotsPtr, Allocator);
            SlotsPtr = null;
            Capacity = 0;
            BucketSize = 0;
            Length = 0;
            FreeIndex = -1;
            LastIndex = 0;
            Allocator = Allocator.Invalid;
            IsCreated = false;
        }

        public TValue this[TKey key]
        {
            [MethodImpl(MethodImplOptions.AggressiveInlining)]
            get
            {
                if (!IsCreated)
                    throw new Exception("UnsafeHashMap is not created");

                int bucketIndex = GetBucketIndex(key);
                int index = BucketsPtr[bucketIndex] - 1;
                while (index >= 0)
                {
                    var slot = &SlotsPtr[index];
                    if (Equals(slot->Key, key))
                        return DataPtr[index];
                    index = slot->NextIndexPlusOne - 1;
                }

                throw new KeyNotFoundException($"get value error! {key} not found in hash map");
            }
            [MethodImpl(MethodImplOptions.AggressiveInlining)]
            set
            {
                if (!IsCreated)
                    throw new Exception("UnsafeHashMap is not created");

                int bucketIndex = GetBucketIndex(key);
                int index = BucketsPtr[bucketIndex] - 1;
                while (index >= 0)
                {
                    var slot = &SlotsPtr[index];
                    if (Equals(slot->Key, key))
                    {
                        DataPtr[index] = value;
                        return;
                    }

                    index = slot->NextIndexPlusOne - 1;
                }

                throw new KeyNotFoundException($"set value error! {key} not found in hash map");
            }
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public ref TValue GetElementRef(TKey key)
        {
            if (!IsCreated)
                throw new Exception("UnsafeHashMap is not created");

            int bucketIndex = GetBucketIndex(key);
            int index = BucketsPtr[bucketIndex] - 1;
            while (index >= 0)
            {
                var slot = &SlotsPtr[index];
                if (Equals(slot->Key, key))
                    return ref DataPtr[index];
                index = slot->NextIndexPlusOne - 1;
            }

            throw new KeyNotFoundException($"get value error! {key} not found in hash map");
        }


        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public TValue* GetElementPtr(TKey key)
        {
            if (!IsCreated)
                throw new Exception("UnsafeHashMap is not created");

            int bucketIndex = GetBucketIndex(key);
            int index = BucketsPtr[bucketIndex] - 1;
            while (index >= 0)
            {
                var slot = &SlotsPtr[index];
                if (Equals(slot->Key, key))
                    return DataPtr + index;
                index = slot->NextIndexPlusOne - 1;
            }

            throw new KeyNotFoundException($"get value ptr error! {key} not found in hash map");
        }

        /// <summary>
        /// 扩容
        /// </summary>
        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public void Expend()
        {
            int newCapacity = GetCapacity(Capacity);
            int newBucketSize = GetCapacity(BucketSize);
            UnsafeHashMap<TKey, TValue> newHashMap =
                new UnsafeHashMap<TKey, TValue>(newCapacity, newBucketSize, Allocator);
            foreach (var item in this)
                newHashMap.Add(item.Key, item.Value);
            UnsafeUtil.Free(DataPtr, Allocator);
            UnsafeUtil.Free(BucketsPtr, Allocator);
            UnsafeUtil.Free(SlotsPtr, Allocator);
            Capacity = newHashMap.Capacity;
            BucketSize = newHashMap.BucketSize;
            DataPtr = newHashMap.DataPtr;
            BucketsPtr = newHashMap.BucketsPtr;
            SlotsPtr = newHashMap.SlotsPtr;
            FreeIndex = -1;
            LastIndex = newHashMap.Length;
        }

        /// <summary>
        /// 扩容
        /// </summary>
        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public void Resize(int newCapacity, int newBucketSize)
        {
            if (Capacity == newCapacity && BucketSize == newBucketSize) return;
            UnsafeHashMap<TKey, TValue> newHashMap =
                new UnsafeHashMap<TKey, TValue>(newCapacity, newBucketSize, Allocator);
            foreach (var item in this)
                newHashMap.Add(item.Key, item.Value);
            UnsafeUtil.Free(DataPtr, Allocator);
            UnsafeUtil.Free(BucketsPtr, Allocator);
            UnsafeUtil.Free(SlotsPtr, Allocator);
            Capacity = newHashMap.Capacity;
            BucketSize = newHashMap.BucketSize;
            DataPtr = newHashMap.DataPtr;
            BucketsPtr = newHashMap.BucketsPtr;
            SlotsPtr = newHashMap.SlotsPtr;
            FreeIndex = -1;
            LastIndex = newHashMap.Length;
        }

        UnsafeSlot* GetSlot(int slotIndex)
        {
            return (UnsafeSlot*)((byte*)SlotsPtr + slotIndex * mSlotSize);
        }

        /// <summary>
        /// 添加键值对
        /// </summary>
        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public bool Add(TKey key, TValue value)
        {
            if (!IsCreated)
                throw new Exception("UnsafeHashMap is not created");

            int bucketIndex = GetBucketIndex(key);
            int v = BucketsPtr[bucketIndex];
            int index = BucketsPtr[bucketIndex] - 1;
            while (index >= 0)
            {
                var slot = &SlotsPtr[index];
                if (Equals(slot->Key, key))
                    return false; //已有key
                index = slot->NextIndexPlusOne - 1;
            }

            if (FreeIndex >= 0)
            {
                index = FreeIndex;
                FreeIndex = (&SlotsPtr[index])->NextIndexPlusOne - 1;
            }
            else
            {
                if (LastIndex == Capacity)
                {
                    Expend();
                    bucketIndex = GetBucketIndex(key);
                }

                index = LastIndex;
                LastIndex++;
            }

            UnsafeSlot* slotPtr = &SlotsPtr[index];
            slotPtr->Key = key;
            slotPtr->NextIndexPlusOne = BucketsPtr[bucketIndex];
            DataPtr[index] = value;
            BucketsPtr[bucketIndex] = index + 1; //IndexPlusOne
            Length++;

            //扩容
            if (1f * Length / BucketSize >= 0.7f)
                Expend();
            return true;
        }

        /// <summary>
        /// 添加键值对集合
        /// </summary>
        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public void AddRange(UnsafeHashMap<TKey, TValue> other)
        {
            if (!IsCreated)
                throw new Exception("UnsafeHashMap is not created");

            if (other.Length == 0) return;
            int newSize = Length + other.Length;
            if (1f * newSize / BucketSize >= 0.7f)
                Resize(GetCapacity(Capacity), GetCapacity(BucketSize));
            foreach (var item in other)
                Add(item.Key, item.Value);
        }

        /// <summary>
        /// 移除键（返回值）
        /// </summary>
        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public bool Remove(TKey key, out TValue value)
        {
            if (!IsCreated)
                throw new Exception("UnsafeHashMap is not created");

            int bucketIndex = GetBucketIndex(key);
            int index = BucketsPtr[bucketIndex] - 1;
            int prevIndex = -1;
            while (index >= 0)
            {
                var slot = &SlotsPtr[index];
                if (Equals(slot->Key, key))
                {
                    if (prevIndex >= 0) //非头节点
                    {
                        if (slot->NextIndexPlusOne > 0) //有子节点
                            (&SlotsPtr[prevIndex])->NextIndexPlusOne = slot->NextIndexPlusOne;
                        else //无子节点
                            (&SlotsPtr[prevIndex])->NextIndexPlusOne = 0;
                    }
                    else //头节点
                        BucketsPtr[bucketIndex] = 0;

                    value = DataPtr[index];
                    DataPtr[index] = default;
                    slot->Key = default;
                    slot->NextIndexPlusOne = FreeIndex + 1; //使用指针记录空格
                    Length--;
                    if (Length == 0)
                    {
                        FreeIndex = -1;
                        LastIndex = 0;
                    }
                    else
                        FreeIndex = index;

                    return true;
                }

                prevIndex = index;
                index = slot->NextIndexPlusOne - 1;
            }

            value = default;
            return false;
        }

        /// <summary>
        /// 移除键
        /// </summary>
        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public bool Remove(TKey key)
        {
            if (!IsCreated)
                throw new Exception("UnsafeHashMap is not created");

            int bucketIndex = GetBucketIndex(key);
            int index = BucketsPtr[bucketIndex] - 1;
            int prevIndex = -1;
            while (index >= 0)
            {
                var slot = &SlotsPtr[index];
                if (Equals(slot->Key, key))
                {
                    if (prevIndex >= 0) //非头节点
                    {
                        if (slot->NextIndexPlusOne > 0) //有子节点
                            (&SlotsPtr[prevIndex])->NextIndexPlusOne = slot->NextIndexPlusOne;
                        else //无子节点
                            (&SlotsPtr[prevIndex])->NextIndexPlusOne = 0;
                    }
                    else //头节点
                        BucketsPtr[bucketIndex] = 0;

                    DataPtr[index] = default;
                    slot->Key = default;
                    slot->NextIndexPlusOne = FreeIndex + 1; //使用指针记录空格
                    Length--;
                    if (Length == 0)
                    {
                        FreeIndex = -1;
                        LastIndex = 0;
                    }
                    else
                        FreeIndex = index;

                    return true;
                }

                prevIndex = index;
                index = slot->NextIndexPlusOne - 1;
            }

            return false;
        }

        /// <summary>
        /// 尝试通过键获取值
        /// </summary>
        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public bool TryGetValue(TKey key, out TValue value)
        {
            if (!IsCreated)
                throw new Exception("UnsafeHashMap is not created");

            int bucketIndex = GetBucketIndex(key);
            int index = BucketsPtr[bucketIndex] - 1;
            while (index >= 0)
            {
                var slot = &SlotsPtr[index];
                if (Equals(slot->Key, key))
                {
                    value = DataPtr[index];
                    return true;
                }

                index = slot->NextIndexPlusOne - 1;
            }

            value = default;
            return false;
        }

        /// <summary>
        /// 清空哈希表
        /// </summary>
        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public void Clear()
        {
            if (Length == 0) return;
            UnsafeUtility.MemClear(DataPtr, Length * mValueSize);
            UnsafeUtility.MemClear(BucketsPtr, BucketSize * mIntSize);
            UnsafeUtility.MemClear(SlotsPtr, Length * mSlotSize);
            Length = 0;
            FreeIndex = -1;
            LastIndex = 0;
        }

        /// <summary>
        /// 检测是否包含指定键
        /// </summary>
        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public bool ContainsKey(TKey key)
        {
            if (!IsCreated)
                throw new Exception("UnsafeHashMap is not created");

            int bucketIndex = GetBucketIndex(key);
            int index = BucketsPtr[bucketIndex] - 1;
            while (index >= 0)
            {
                var slot = &SlotsPtr[index];
                if (Equals(slot->Key, key))
                    return true;
                index = slot->NextIndexPlusOne - 1;
            }

            return false;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        int GetCapacity(int currentSize)
        {
            return currentSize * 2;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        int GetBucketIndex(TKey key)
        {
            int hashCode = key.GetHashCode();
            return (hashCode > 0 ? hashCode : -hashCode) % BucketSize;
        }

        public struct UnsafeSlot
        {
            public TKey Key;
            public int NextIndexPlusOne;
        }

        IEnumerator<KeyValuePair<TKey, TValue>> IEnumerable<KeyValuePair<TKey, TValue>>.GetEnumerator()
        {
            return GetEnumerator();
        }

        IEnumerator IEnumerable.GetEnumerator()
        {
            return GetEnumerator();
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public Enumerator GetEnumerator()
        {
            return new Enumerator(this);
        }

        public struct Enumerator : IEnumerator<KeyValuePair<TKey, TValue>>, IDisposable
        {
            UnsafeHashMap<TKey, TValue> mHashMap;
            int mIndex;
            int mCurrentIndex; //当前索引
            int mBucketIndex; //在bucket中的索引
            int mNodeIndex; //在链表中的索引
            bool mIsLeaf; //是否为链表的叶子节点

            public KeyValuePair<TKey, TValue> Current
            {
                [MethodImpl(MethodImplOptions.AggressiveInlining)]
                get
                {
                    TKey key = (&mHashMap.SlotsPtr[mCurrentIndex])->Key;
                    TValue value = mHashMap.DataPtr[mCurrentIndex];
                    return new KeyValuePair<TKey, TValue>(key, value);
                }
            }

            object IEnumerator.Current => Current;

            public Enumerator(UnsafeHashMap<TKey, TValue> hashMap)
            {
                mHashMap = hashMap;
                mIndex = -1;
                mCurrentIndex = -1;
                mBucketIndex = -1;
                mNodeIndex = -1;
                mIsLeaf = true;
            }

            public bool MoveNext()
            {
                mIndex++;
                if (mIndex >= mHashMap.Length)
                    return false;
                mNodeIndex = mIsLeaf ? 0 : mNodeIndex + 1;
                for (int i = mBucketIndex + (mIsLeaf ? 1 : 0), bucketSize = mHashMap.BucketSize; i < bucketSize; i++)
                {
                    mCurrentIndex = mHashMap.BucketsPtr[i] - 1;
                    if (mCurrentIndex < 0) continue;
                    for (int j = 0; j < mNodeIndex; j++)
                    {
                        mCurrentIndex = (&mHashMap.SlotsPtr[mCurrentIndex])->NextIndexPlusOne - 1;
                        if (mCurrentIndex < 0) break;
                    }

                    if (mCurrentIndex < 0)
                    {
                        mNodeIndex = 0;
                        continue;
                    }

                    mBucketIndex = i;
                    mIsLeaf = (&mHashMap.SlotsPtr[mCurrentIndex])->NextIndexPlusOne == 0;
                    break;
                }

                return true;
            }

            public void Reset()
            {
                mIndex = -1;
                mCurrentIndex = -1;
                mBucketIndex = -1;
                mNodeIndex = -1;
                mIsLeaf = true;
            }

            public void Dispose()
            {
            }
        }
    }
}